{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b076bd1a-b236-4fbc-953d-8295b25122ae",
   "metadata": {},
   "source": [
    "# 👖 Variational Autoencoders - Fashion-MNIST"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9235cbd1-f136-411c-88d9-f69f270c0b96",
   "metadata": {},
   "source": [
    "In this notebook, we'll walk through the steps required to train your own autoencoder on the fashion MNIST dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ab8071b-6381-4147-8755-ea2ee48e8ab3",
   "metadata": {},
   "source": [
    "The code has been adapted from the excellent [VAE tutorial](https://keras.io/examples/generative/vae/) created by Francois Chollet, available on the Keras website."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84acc7be-6764-4668-b2bb-178f63deeed3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow.keras.backend as K\n",
    "from tensorflow.keras import (\n",
    "    layers,\n",
    "    models,\n",
    "    datasets,\n",
    "    callbacks,\n",
    "    losses,\n",
    "    optimizers,\n",
    "    metrics,\n",
    ")\n",
    "\n",
    "from scipy.stats import norm\n",
    "\n",
    "from notebooks.utils import display"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "339e6268-ebd7-4feb-86db-1fe7abccdbe5",
   "metadata": {},
   "source": [
    "## 0. Parameters <a name=\"parameters\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b2ee6ce-129f-4833-b0c5-fa567381c4e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "IMAGE_SIZE = 32\n",
    "BATCH_SIZE = 100\n",
    "VALIDATION_SPLIT = 0.2\n",
    "EMBEDDING_DIM = 2\n",
    "EPOCHS = 5\n",
    "BETA = 500"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7716fac-0010-49b0-b98e-53be2259edde",
   "metadata": {},
   "source": [
    "## 1. Prepare the data <a name=\"prepare\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a73e5a4-1638-411c-8d3c-29f823424458",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the data\n",
    "(x_train, y_train), (x_test, y_test) = datasets.fashion_mnist.load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebae2f0d-59fd-4796-841f-7213eae638de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preprocess the data\n",
    "\n",
    "\n",
    "def preprocess(imgs):\n",
    "    \"\"\"\n",
    "    Normalize and reshape the images\n",
    "    \"\"\"\n",
    "    imgs = imgs.astype(\"float32\") / 255.0\n",
    "    imgs = np.pad(imgs, ((0, 0), (2, 2), (2, 2)), constant_values=0.0)\n",
    "    imgs = np.expand_dims(imgs, -1)\n",
    "    return imgs\n",
    "\n",
    "\n",
    "x_train = preprocess(x_train)\n",
    "x_test = preprocess(x_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa53709f-7f3f-483b-9db8-2e5f9b9942c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show some items of clothing from the training set\n",
    "display(x_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aff50401-3abe-4c10-bba8-b35bc13ad7d5",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 2. Build the variational autoencoder <a name=\"build\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a0625b6-3c19-478b-84f9-5c2b5c2b74b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sampling(layers.Layer):\n",
    "    def call(self, inputs):\n",
    "        z_mean, z_log_var = inputs\n",
    "        batch = tf.shape(z_mean)[0]\n",
    "        dim = tf.shape(z_mean)[1]\n",
    "        epsilon = K.random_normal(shape=(batch, dim))\n",
    "        return z_mean + tf.exp(0.5 * z_log_var) * epsilon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "086e2584-c60d-4990-89f4-2092c44e023e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Encoder\n",
    "encoder_input = layers.Input(\n",
    "    shape=(IMAGE_SIZE, IMAGE_SIZE, 1), name=\"encoder_input\"\n",
    ")\n",
    "x = layers.Conv2D(32, (3, 3), strides=2, activation=\"relu\", padding=\"same\")(\n",
    "    encoder_input\n",
    ")\n",
    "x = layers.Conv2D(64, (3, 3), strides=2, activation=\"relu\", padding=\"same\")(x)\n",
    "x = layers.Conv2D(128, (3, 3), strides=2, activation=\"relu\", padding=\"same\")(x)\n",
    "shape_before_flattening = K.int_shape(x)[1:]  # the decoder will need this!\n",
    "\n",
    "x = layers.Flatten()(x)\n",
    "z_mean = layers.Dense(EMBEDDING_DIM, name=\"z_mean\")(x)\n",
    "z_log_var = layers.Dense(EMBEDDING_DIM, name=\"z_log_var\")(x)\n",
    "z = Sampling()([z_mean, z_log_var])\n",
    "\n",
    "encoder = models.Model(encoder_input, [z_mean, z_log_var, z], name=\"encoder\")\n",
    "encoder.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "928f56bd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c409e63-1aea-42e2-8324-c3e2a12073ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Decoder\n",
    "decoder_input = layers.Input(shape=(EMBEDDING_DIM,), name=\"decoder_input\")\n",
    "x = layers.Dense(np.prod(shape_before_flattening))(decoder_input)\n",
    "x = layers.Reshape(shape_before_flattening)(x)\n",
    "x = layers.Conv2DTranspose(\n",
    "    128, (3, 3), strides=2, activation=\"relu\", padding=\"same\"\n",
    ")(x)\n",
    "x = layers.Conv2DTranspose(\n",
    "    64, (3, 3), strides=2, activation=\"relu\", padding=\"same\"\n",
    ")(x)\n",
    "x = layers.Conv2DTranspose(\n",
    "    32, (3, 3), strides=2, activation=\"relu\", padding=\"same\"\n",
    ")(x)\n",
    "decoder_output = layers.Conv2D(\n",
    "    1,\n",
    "    (3, 3),\n",
    "    strides=1,\n",
    "    activation=\"sigmoid\",\n",
    "    padding=\"same\",\n",
    "    name=\"decoder_output\",\n",
    ")(x)\n",
    "\n",
    "decoder = models.Model(decoder_input, decoder_output)\n",
    "decoder.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65b79c40-a75a-4352-935e-d644632b4d9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class VAE(models.Model):\n",
    "    def __init__(self, encoder, decoder, **kwargs):\n",
    "        super(VAE, self).__init__(**kwargs)\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.total_loss_tracker = metrics.Mean(name=\"total_loss\")\n",
    "        self.reconstruction_loss_tracker = metrics.Mean(\n",
    "            name=\"reconstruction_loss\"\n",
    "        )\n",
    "        self.kl_loss_tracker = metrics.Mean(name=\"kl_loss\")\n",
    "\n",
    "    @property\n",
    "    def metrics(self):\n",
    "        return [\n",
    "            self.total_loss_tracker,\n",
    "            self.reconstruction_loss_tracker,\n",
    "            self.kl_loss_tracker,\n",
    "        ]\n",
    "\n",
    "    def call(self, inputs):\n",
    "        \"\"\"Call the model on a particular input.\"\"\"\n",
    "        z_mean, z_log_var, z = encoder(inputs)\n",
    "        reconstruction = decoder(z)\n",
    "        return z_mean, z_log_var, reconstruction\n",
    "\n",
    "    def train_step(self, data):\n",
    "        \"\"\"Step run during training.\"\"\"\n",
    "        with tf.GradientTape() as tape:\n",
    "            z_mean, z_log_var, reconstruction = self(data)\n",
    "            reconstruction_loss = tf.reduce_mean(\n",
    "                BETA\n",
    "                * losses.binary_crossentropy(\n",
    "                    data, reconstruction, axis=(1, 2, 3)\n",
    "                )\n",
    "            )\n",
    "            kl_loss = tf.reduce_mean(\n",
    "                tf.reduce_sum(\n",
    "                    -0.5\n",
    "                    * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)),\n",
    "                    axis=1,\n",
    "                )\n",
    "            )\n",
    "            total_loss = reconstruction_loss + kl_loss\n",
    "\n",
    "        grads = tape.gradient(total_loss, self.trainable_weights)\n",
    "        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))\n",
    "\n",
    "        self.total_loss_tracker.update_state(total_loss)\n",
    "        self.reconstruction_loss_tracker.update_state(reconstruction_loss)\n",
    "        self.kl_loss_tracker.update_state(kl_loss)\n",
    "\n",
    "        return {m.name: m.result() for m in self.metrics}\n",
    "\n",
    "    def test_step(self, data):\n",
    "        \"\"\"Step run during validation.\"\"\"\n",
    "        if isinstance(data, tuple):\n",
    "            data = data[0]\n",
    "\n",
    "        z_mean, z_log_var, reconstruction = self(data)\n",
    "        reconstruction_loss = tf.reduce_mean(\n",
    "            BETA\n",
    "            * losses.binary_crossentropy(data, reconstruction, axis=(1, 2, 3))\n",
    "        )\n",
    "        kl_loss = tf.reduce_mean(\n",
    "            tf.reduce_sum(\n",
    "                -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)),\n",
    "                axis=1,\n",
    "            )\n",
    "        )\n",
    "        total_loss = reconstruction_loss + kl_loss\n",
    "\n",
    "        return {\n",
    "            \"loss\": total_loss,\n",
    "            \"reconstruction_loss\": reconstruction_loss,\n",
    "            \"kl_loss\": kl_loss,\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edf2f892-9209-42ee-b251-1e7604df5335",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a variational autoencoder\n",
    "vae = VAE(encoder, decoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35b14665-4359-447b-be58-3fd58ba69084",
   "metadata": {},
   "source": [
    "## 3. Train the variational autoencoder <a name=\"train\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b429fdad-ea9c-45a2-a556-eb950d793824",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compile the variational autoencoder\n",
    "optimizer = optimizers.Adam(learning_rate=0.0005)\n",
    "vae.compile(optimizer=optimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c525e44b-b3bb-489c-9d35-fcfe3e714e6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a model save checkpoint\n",
    "model_checkpoint_callback = callbacks.ModelCheckpoint(\n",
    "    filepath=\"./checkpoint\",\n",
    "    save_weights_only=False,\n",
    "    save_freq=\"epoch\",\n",
    "    monitor=\"loss\",\n",
    "    mode=\"min\",\n",
    "    save_best_only=True,\n",
    "    verbose=0,\n",
    ")\n",
    "tensorboard_callback = callbacks.TensorBoard(log_dir=\"./logs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3c497b7-fa40-48df-b2bf-541239cc9400",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "vae.fit(\n",
    "    x_train,\n",
    "    epochs=EPOCHS,\n",
    "    batch_size=BATCH_SIZE,\n",
    "    shuffle=True,\n",
    "    validation_data=(x_test, x_test),\n",
    "    callbacks=[model_checkpoint_callback, tensorboard_callback],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "028138af-d3a5-4134-b980-d3a8a703e70f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the final models\n",
    "vae.save(\"./models/vae\")\n",
    "encoder.save(\"./models/encoder\")\n",
    "decoder.save(\"./models/decoder\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07ab76a1-c4b8-44e3-9c0f-00622fa82277",
   "metadata": {},
   "source": [
    "## 3. Reconstruct using the variational autoencoder <a name=\"reconstruct\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d33cd456-ad3c-4086-be4c-c937e8e6cb90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select a subset of the test set\n",
    "n_to_predict = 5000\n",
    "example_images = x_test[:n_to_predict]\n",
    "example_labels = y_test[:n_to_predict]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af8c8327-fda5-40d7-8909-9883c47b1a3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create autoencoder predictions and display\n",
    "z_mean, z_log_var, reconstructions = vae.predict(example_images)\n",
    "print(\"Example real clothing items\")\n",
    "display(example_images)\n",
    "print(\"Reconstructions\")\n",
    "display(reconstructions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caf7d66e-b194-4372-bcc2-cdda026f8593",
   "metadata": {},
   "source": [
    "## 4. Embed using the encoder <a name=\"encode\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aac4bca-15fa-4679-b790-1347854c22b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Encode the example images\n",
    "z_mean, z_var, z = encoder.predict(example_images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6836c5d-a7f3-4a98-ba1f-3d204a8654c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some examples of the embeddings\n",
    "print(z[:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02736f3a-2446-4e4c-a8c1-f1eb34729f03",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the encoded points in 2D space\n",
    "figsize = 8\n",
    "\n",
    "plt.figure(figsize=(figsize, figsize))\n",
    "plt.scatter(z[:, 0], z[:, 1], c=\"black\", alpha=0.5, s=3)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8fb22e1-f73f-4b97-8a8a-787c1c5e605a",
   "metadata": {},
   "source": [
    "## 5. Generate using the decoder <a name=\"decode\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8cbb2e9-4ba6-4332-b5cf-c509472b4d39",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample some points in the latent space, from the standard normal distribution\n",
    "grid_width, grid_height = (6, 3)\n",
    "z_sample = np.random.normal(size=(grid_width * grid_height, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd914b15-f9f5-4b59-934d-52ae4c7917cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Decode the sampled points\n",
    "reconstructions = decoder.predict(z_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc563362-4683-4065-a4f5-f2bbe21ce62e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert original embeddings and sampled embeddings to p-values\n",
    "p = norm.cdf(z)\n",
    "p_sample = norm.cdf(z_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0159b056-675e-4f3d-ab1a-eab40cfa9c46",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Draw a plot of...\n",
    "figsize = 8\n",
    "plt.figure(figsize=(figsize, figsize))\n",
    "\n",
    "# ... the original embeddings ...\n",
    "plt.scatter(z[:, 0], z[:, 1], c=\"black\", alpha=0.5, s=2)\n",
    "\n",
    "# ... and the newly generated points in the latent space\n",
    "plt.scatter(z_sample[:, 0], z_sample[:, 1], c=\"#00B0F0\", alpha=1, s=40)\n",
    "plt.show()\n",
    "\n",
    "# Add underneath a grid of the decoded images\n",
    "fig = plt.figure(figsize=(figsize, grid_height * 2))\n",
    "fig.subplots_adjust(hspace=0.4, wspace=0.4)\n",
    "\n",
    "for i in range(grid_width * grid_height):\n",
    "    ax = fig.add_subplot(grid_height, grid_width, i + 1)\n",
    "    ax.axis(\"off\")\n",
    "    ax.text(\n",
    "        0.5,\n",
    "        -0.35,\n",
    "        str(np.round(z_sample[i, :], 1)),\n",
    "        fontsize=10,\n",
    "        ha=\"center\",\n",
    "        transform=ax.transAxes,\n",
    "    )\n",
    "    ax.imshow(reconstructions[i, :, :], cmap=\"Greys\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a70a79b8-b109-4a47-8364-f0edc15a0e7f",
   "metadata": {},
   "source": [
    "## 6. Explore the latent space <a name=\"explore\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9547247c-0e2a-4ac5-9490-4a1346a3f527",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Colour the embeddings by their label (clothing type - see table)\n",
    "figsize = 8\n",
    "fig = plt.figure(figsize=(figsize * 2, figsize))\n",
    "ax = fig.add_subplot(1, 2, 1)\n",
    "plot_1 = ax.scatter(\n",
    "    z[:, 0], z[:, 1], cmap=\"rainbow\", c=example_labels, alpha=0.8, s=3\n",
    ")\n",
    "plt.colorbar(plot_1)\n",
    "ax = fig.add_subplot(1, 2, 2)\n",
    "plot_2 = ax.scatter(\n",
    "    p[:, 0], p[:, 1], cmap=\"rainbow\", c=example_labels, alpha=0.8, s=3\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e176e97-8ee6-491a-9f93-92e9e526a2c7",
   "metadata": {},
   "source": [
    "| ID | Clothing Label |\n",
    "| :- | :- |\n",
    "| 0 | T-shirt/top |\n",
    "| 1 | Trouser |\n",
    "| 2 | Pullover |\n",
    "| 3 | Dress |\n",
    "| 4 | Coat |\n",
    "| 5 | Sandal |\n",
    "| 6 | Shirt |\n",
    "| 7 | Sneaker |\n",
    "| 8 | Bag |\n",
    "| 9 | Ankle boot |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "276fdfa2-b5e9-41fe-b13d-fde00e3edff4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Colour the embeddings by their label (clothing type - see table)\n",
    "figsize = 12\n",
    "grid_size = 15\n",
    "plt.figure(figsize=(figsize, figsize))\n",
    "plt.scatter(\n",
    "    p[:, 0], p[:, 1], cmap=\"rainbow\", c=example_labels, alpha=0.8, s=300\n",
    ")\n",
    "plt.colorbar()\n",
    "\n",
    "x = norm.ppf(np.linspace(0, 1, grid_size))\n",
    "y = norm.ppf(np.linspace(1, 0, grid_size))\n",
    "xv, yv = np.meshgrid(x, y)\n",
    "xv = xv.flatten()\n",
    "yv = yv.flatten()\n",
    "grid = np.array(list(zip(xv, yv)))\n",
    "\n",
    "reconstructions = decoder.predict(grid)\n",
    "# plt.scatter(grid[:, 0], grid[:, 1], c=\"black\", alpha=1, s=10)\n",
    "plt.show()\n",
    "\n",
    "fig = plt.figure(figsize=(figsize, figsize))\n",
    "fig.subplots_adjust(hspace=0.4, wspace=0.4)\n",
    "for i in range(grid_size**2):\n",
    "    ax = fig.add_subplot(grid_size, grid_size, i + 1)\n",
    "    ax.axis(\"off\")\n",
    "    ax.imshow(reconstructions[i, :, :], cmap=\"Greys\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
